#
#  GoogleFindMyTools - A set of tools to interact with the Google Find My API
#  Copyright © 2024 Leon Böttger. All rights reserved.
#
from Cryptodome.Cipher import AES
from ecdsa import SECP160r1

from example_data_provider import get_example_data

# Constants
K = 10
ROTATION_PERIOD = 1024  # 2^K seconds

def generate_eid(identity_key: bytes, timestamp: int) -> bytes:
    # Calculate r
    r = calculate_r(identity_key, timestamp)

    # Compute R = r * G
    curve = SECP160r1
    R = r * curve.generator

    # Return the x coordinate of R as the EID
    return R.x().to_bytes(20, 'big')


def calculate_r(identity_key: bytes, timestamp: int):
    # ts_bytes is the timestamp in bytes, but the least K significant bits are set to 0
    ts_bytes = get_masked_timestamp(timestamp, K)
    identity_key_bytes = identity_key

    # A random is generated by AES-ECB-256 encrypting the following data structure with the ephemeral identity key:
    data = bytearray(32)
    data[0:11] = b'\xFF' * 11
    data[11] = K
    data[12:16] = ts_bytes
    data[16:27] = b'\x00' * 11
    data[27] = K
    data[28:32] = ts_bytes

    # AES-ECB-256 encryption
    cipher = AES.new(identity_key_bytes, AES.MODE_ECB)
    r_dash = cipher.encrypt(bytes(data))

    # Convert r' to an integer
    r_dash_int = int.from_bytes(r_dash, byteorder='big', signed=False)

    # SECP160R1 parameters
    curve = SECP160r1
    n = curve.order

    # r' is now projected to the finite field Fp by calculating r = r' mod n
    return (r_dash_int % n)


def get_masked_timestamp(timestamp: int, K: int):
    # Create a bitmask that has all bits set except for the K least significant bits
    mask = ~((1 << K) - 1)

    # Zero out the K least significant bits
    timestamp &= mask

    # Convert back to a byte array with the same length as the original
    return timestamp.to_bytes(4, byteorder='big')


if __name__ == '__main__':

    sample_identity_key = get_example_data("sample_identity_key")

    # Generate EIDs
    for i in range(1000):
        timestamp = i * ROTATION_PERIOD
        eid = generate_eid(sample_identity_key, timestamp)
        print(f"{timestamp}: {eid.hex()}")